//
// Created by liuhuan on 2021/5/22.
//

#include <gtest/gtest.h>
#include "random"
#include <string>
#include "OPE.h"
#include "ValueRange.h"
#include "Coins.h"
#include "Hgd.h"

using namespace std;
using std::string;
using std::random_device;
using std::default_random_engine;

OPE::OPE() {
    this->key = genRandKey(DEFAULT_LENGTH);
    inRange = ValueRange(DEFAULT_IN_RANGE_START, DEFAULT_IN_RANGE_END);
    outRange = ValueRange(DEFAULT_OUT_RANGE_START, DEFAULT_OUT_RANGE_END);
    if (inRange.size() > outRange.size()) {
        throw range_error("Invalid range");
    }
}

OPE::OPE(string key) {
    this->key = key;
    inRange = ValueRange(DEFAULT_IN_RANGE_START, DEFAULT_IN_RANGE_END);
    outRange = ValueRange(DEFAULT_OUT_RANGE_START, DEFAULT_OUT_RANGE_END);
    if (inRange.size() > outRange.size()) {
        throw range_error("Invalid range");
    }
}

OPE::OPE(string key, ValueRange inRange, ValueRange outRange) {
    this->key = key;
    this->inRange = inRange;
    this->outRange = outRange;
    if (inRange.size() > outRange.size()) {
        throw range_error("Invalid range");
    }
}

uint OPE::encrypt(uint plaintext) {
    if (!this->inRange.contains(plaintext)) {
        throw range_error("Plaintext is not within the input range");
    }
    return encryptRecursive(plaintext, this->inRange, this->outRange);
}

/**
 * 保序加密 的具体加密过程
 */
uint OPE::encryptRecursive(uint plaintext, ValueRange inRange, ValueRange outRange) {
    uint inSize = inRange.size();           // M
    uint outSize = outRange.size();         // N
    uint inEdge = inRange.start - 1;        // d
    uint outEdge = outRange.start - 1;      // r
    uint mid = outEdge + uint(ceil(outSize / 2));   // y    math.ceil 向上取整
    assert(inSize <= outSize);
    if (inRange.size() == 1) {          // 明文空间只有一个，即 plaintext
        Coins coins = Coins(key, plaintext);
        uint ciphertext = sampleUniform(outRange, coins);// 返回最终的 抽样结果
        return ciphertext;
    }
    Coins coins = Coins(key, mid);
    uint x = sampleHgd(inRange, outRange, mid, coins);  // 中间过程的 抽样值

    if (plaintext <= x) {
        inRange = ValueRange(inEdge + 1, x);
        outRange = ValueRange(outEdge + 1, mid);
    } else {
        inRange = ValueRange(x + 1, inEdge + inSize);
        outRange = ValueRange(mid + 1, outEdge + outSize);
    }
    return encryptRecursive(plaintext, inRange, outRange);
}

string OPE::genRandKey(int length) {
    char tmp;
    string buffer;

    random_device rd;
    default_random_engine random(rd());

    for (int i = 0; i < length; i++) {
        tmp = random() % 36;
        if (tmp < 10) {
            tmp += '0';
        } else {
            tmp -= 10;
            tmp += 'A';
        }
        buffer += tmp;
    }
    return buffer;
}

/**
 * 一个抽样函数
 */
uint OPE::sampleHgd(ValueRange inRange, ValueRange outRage, uint nSample, Coins seedCoins) {
    uint inSize = inRange.size();
    uint outSize = outRage.size();
    assert(inSize > 0 and outSize > 0);
    assert(inSize <= outSize);
    assert(outRage.contains(nSample));

    uint nSampleIndex = nSample - outRage.start + 1;
    if (inSize == outSize) {
        return inRange.start + nSampleIndex - 1;
    }

    uint inSampleNum = Hgd::rhyper(nSampleIndex, inSize, outSize - inSize, seedCoins);
    if (inSampleNum == 0) {
        return inRange.start;
    } else {
        uint inSample = inRange.start + inSampleNum -1;
        assert(inRange.contains(inSample));
        return inSample;
    }
}

/**
 * 一个抽样函数
 */
uint OPE::sampleUniform(ValueRange inRange, Coins coins) {
    ValueRange curRange = ValueRange(inRange);
    assert(curRange.size() != 0);
    while (curRange.size() > 0) {
        uint mid = (curRange.start + curRange.end) / 2;
        bool bit = coins.next();
        if (bit == false) {
            curRange.end = mid;
        } else if (bit == true) {
            curRange.start = mid +1;
        } else {
            throw invalid_argument("Unexpected bit value");
        }
    }
    assert(curRange.size() == 0);
    return curRange.start;
}


/**
 * 使用 gtest 框架测试代码，可以忽略
 */
TEST(testOPE, testOrderGuarantees) {
    uint values[9] = {1, 2, 10, 28, 42, 1000, 2000, 4001};
    uint result[9];
    string key = "key";
    OPE cipher = OPE(key);
    for (int i = 0; i < 8; ++i) {
        result[i] = cipher.encrypt(values[i]);
    }
    for (int i = 0; i < 7; ++i) {
        EXPECT_LT(result[i], result[i+1]);
    }
}

TEST(testOPE, testOpeDeterministic) {
    uint values[5] = {1, 314, 1337, 1338, 10000};
    uint encryptedValuesFirst[5];
    uint encryptedValuesSecond[5];
    OPE cipher = OPE("key-la-la");
    for (int i = 0; i < 5; ++i) {
        encryptedValuesFirst[i] = cipher.encrypt(values[i]);
    }
    for (int i = 0; i < 5; ++i) {
        encryptedValuesSecond[i] = cipher.encrypt(values[i]);
    }
    for (int i = 0; i < 5; ++i) {
        EXPECT_EQ(encryptedValuesFirst[i], encryptedValuesSecond[i]);
    }
}